%Code to obtain VCV matrix of simulated moments. 
%inputs:
%1. # of iterations (Nreps) 
%2. # of simualted moments (Nmoms)
%3. cells holding estimated policy functions (optfunc & optfunc_frz)
%4. vcv_params: z_age_min,z_age_max,N_obs,mu,rho_sigma_v,f_N,aT,aS
%5. relative scale of treatment effects compared to employment rate
%(scaler)
%outputs:
%1. matrix of iteratively computed simulated moments (Sm)
%2. matrix of the VCV of simulated moments (W)

function [Sm,W_matrix] = vcv_matrix(Nreps,Nmoms,optfunc,optfunc_frz,vcv_params,scaler)

    %Define some params
    z_age_min = 56; %minimum freeze age
    z_age_max = 64; %max freeze age
    N_obs     = 5000; %number of simulated individuals
    mu        = 0; %mean of AR(1) -- zero
    rho       = vcv_params(1); %step 1 estimate of persistence of f (AR1)
    sigma_v   = vcv_params(2); %step 1 estimate of variance of  v (AR1)
    f_N       = 6; %# of discrete levels of f
    aT        = 81; %max age
    aS        = 50; %min age
    r         = .029; %interest rate
    pi        = .028;          %inflation
    
    %-------------------------------------------------------------------------%
    % Asset grids (A = non-pension; W = DC pension)
    %-------------------------------------------------------------------------%
    %A
    A_gridN = 20;
    split = 5;
    A_mins = [0;50000;150000;500000;1000000;2000000];
    As=zeros(A_gridN/split,split);
    for i = 1:5
        step = (A_mins(i+1)-A_mins(i))/(A_gridN/split);
        As(:,i) = A_mins(i):step:(A_mins(i+1)-step);
    end
    A = reshape(As,[A_gridN,1]);

    %W
    W_gridN = 20;
    split = 5;
    W_mins = [0;25000;75000;250000;500000;1500000];
    Ws=zeros(W_gridN/split,split);
    for i = 1:5
        step = (W_mins(i+1)-W_mins(i))/(W_gridN/split);
        Ws(:,i) = W_mins(i):step:(W_mins(i+1)-step);
    end
    W = reshape(Ws,[W_gridN,1]);    

    %Pre-create a matrix to hold simulated moments
    Sm = zeros(Nreps,Nmoms);
    seeds = 1:1:Nreps;

        parfor row = 1:Nreps

                %Set seed
                rng(seeds(row),'twister');
                s = rng;
                %-----------------------------------------------------%
                %Simulate shocks and solve for labor supply and saving
                %-----------------------------------------------------%

                %Cell arrays to hold simulated paths 
                LsimC = cell(z_age_max-z_age_min+1,1);
                AsimC = cell(z_age_max-z_age_min+1,1);
                WsimC = cell(z_age_max-z_age_min+1,1);
                LsimT = cell(z_age_max-z_age_min+1,1);
                AsimT = cell(z_age_max-z_age_min+1,1);
                WsimT = cell(z_age_max-z_age_min+1,1);

                %Nodes of discrete AR(1) process -- need max and min of these to
                %assign values out of bounds in the simulation
                [f,~] = rouwenhorst(f_N,mu,rho,sigma_v);
                hor = 12; %horizon

                %Simulated observations per age (prop to LFP in HRS)
                age_shares = csvread('age_shares.csv',1,1);
                obs = round(N_obs.*age_shares(:,1));

                %Disutility shock process
                D = cell(9,1);
                for a = 1:9
                    %simulate AR(1) shocks. 1000 period burn-in... 
                    f_sim0 = zeros(1000+(hor+6),obs(a));
                    rng(s);
                    v = normrnd(0,sigma_v,[length(f_sim0),obs(a)]);
                    for i = 2:length(f_sim0)
                        f_sim0(i,:) = (1-rho)*mu+rho*f_sim0(i-1,:)+v(i,:);
                    end
                    %Trim away burn-in periods
                    f_sim0 = f_sim0(end-(hor+6-1):end,:);
                    D{a,1} = f_sim0;
                end

                %Initial assets --moments from HRS
                A0 = cell(9,1);
                W0 = cell(9,1);
                A_params = csvread('A_params.csv',1,0);
                J_params = csvread('J_params.csv',1,0);
                for a = 1:9
                    rng(s);
                    %Number of individuals with 0 DC wealth
                    frac0obs = round(A_params(a,4)*obs(a)); 
                    %Non-DC wealth for individuals with 0 DC wealth
                    A_temp0 = lognrnd(A_params(a,2),A_params(a,3),[frac0obs,1]);
                    %DC wealth for individuals with no DC balances = 0
                    W_temp0 = zeros(frac0obs,1);
                    %Joint wealth distribution for individuals with >0 DC wealth
                    mvnmu = J_params(a,2:3);
                    mvnsig = [(J_params(a,4))^2 J_params(a,6);J_params(a,6) (J_params(a,5))^2];
                    J_wealth = exp(mvnrnd(mvnmu,mvnsig,(obs(a)-frac0obs)));
                    W_temp1 = J_wealth(:,1);
                    A_temp1 = J_wealth(:,2);
                    A0{a,1} = [A_temp0' A_temp1'];
                    W0{a,1} = [W_temp0' W_temp1'];
                end


                %use policy functions to obtain optimal labor supply and asset
                %accumulation for each simulated individual in each period 
                %-------------------------------------------------------------------------%
                %Control group
                %-------------------------------------------------------------------------%

                for p = 1:9 
                    L_sim = zeros(hor+6,obs(p));
                    A_sim = zeros(hor+6,obs(p));
                    W_sim = zeros(hor+6,obs(p));
                    A_sim(1,:) = A0{p,1};
                    W_sim(1,:) = W0{p,1};

                    %Set disutility distribution
                    f_sim = D{p,1};

                    %Employment choice for t = -5 outside the loop (t = period
                    %relative to freeze) a = 1 when t = -5
                    for i = 1:obs(p)
                        a=1;    
                        Afunc_w = optfunc{(aT-aS)-(p+a-1),2};
                        Wfunc_w = optfunc{(aT-aS)-(p+a-1),3};      
                        Afunc_r1 = optfunc{(aT-aS)-(p+a-1),5};
                        Afunc_r  = Afunc_r1{p+a};
                        Wfunc_r1 = optfunc{(aT-aS)-(p+a-1),6};  
                        Wfunc_r  = Wfunc_r1{p+a};
                        Lfunc = optfunc{(aT-aS)-(p+a-1),8};
                        %Put extreme draws of A,W,f onto the grid 
                        %A
                        if A_sim(a,i)>max(A)
                            A_sim(a,i) = max(A);
                        end
                        %W
                        if W_sim(a,i)>max(W)
                            W_sim(a,i) = max(W);
                        end	
                        %f                
                        if f_sim(a,i)<min(f) 
                            f_sim(a,i) = min(f);
                        elseif f_sim(a,i)>max(f)
                            f_sim(a,i) = max(f);
                        end      

                        %Simulated values 
                        L_sim(a,i) = round(lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Lfunc,A_sim(a,i),W_sim(a,i),f_sim(a,i)));

                        if L_sim(a,i) == 1
                            A_sim(a+1,i) = lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Afunc_w,A_sim(a,i),W_sim(a,i),f_sim(a,i));
                            W_sim(a+1,i) = lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Wfunc_w,A_sim(a,i),W_sim(a,i),f_sim(a,i));
                        elseif L_sim(a,i) == 0 
                            A_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Afunc_r,A_sim(a,i),W_sim(a,i));
                            W_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Wfunc_r,A_sim(a,i),W_sim(a,i));   
                        end  
                    end

                    for a = 2:hor+5 
                       for i = 1:obs(p)

                            Afunc_w = optfunc{(aT-aS)-(p+a-1),2};
                            Wfunc_w = optfunc{(aT-aS)-(p+a-1),3};      
                            Afunc_r1 = optfunc{(aT-aS)-(p+a-1),5};
                            Afunc_r  = Afunc_r1{p+a};
                            Wfunc_r1 = optfunc{(aT-aS)-(p+a-1),6};  
                            Wfunc_r  = Wfunc_r1{p+a};                  
                            Lfunc = optfunc{(aT-aS)-(p+a-1),8};

                            %Put extreme draws of A,W,f onto the grid 
                            %A
                            if A_sim(a,i)>max(A)
                                A_sim(a,i) = max(A);
                            end
                            %W
                            if W_sim(a,i)>max(W)
                                W_sim(a,i) = max(W);
                            end	
                            %f                
                            if f_sim(a,i)<min(f) 
                                f_sim(a,i) = min(f);
                            elseif f_sim(a,i)>max(f)
                                f_sim(a,i) = max(f);
                            end        

                            %Enforce self-absorbing retirement (value functions impose it)
                            if  L_sim(a-1,i) == 1                 
                                L_sim(a,i) = round(lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Lfunc,A_sim(a,i),W_sim(a,i),f_sim(a,i)));                   
                                    if L_sim(a,i) == 1
                                        A_sim(a+1,i) = lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Afunc_w,A_sim(a,i),W_sim(a,i),f_sim(a,i));
                                        W_sim(a+1,i) = lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Wfunc_w,A_sim(a,i),W_sim(a,i),f_sim(a,i));
                                    elseif L_sim(a,i) == 0 
                                        A_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Afunc_r,A_sim(a,i),W_sim(a,i));
                                        W_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Wfunc_r,A_sim(a,i),W_sim(a,i));   
                                    end                    
                            elseif L_sim(a-1,i) == 0
                                L_sim(a,i) = 0;
                                A_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Afunc_r,A_sim(a,i),W_sim(a,i));
                                W_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Wfunc_r,A_sim(a,i),W_sim(a,i));
                            end

                        end
                    end

                    %Final period employment choice (A,W already determined)
                    a = hor+6;
                    for i = 1:obs(p)

                            Lfunc = optfunc{(aT-aS)-(p+a-1),8};

                            %Put extreme draws of A,W,f onto the grid 
                            %A
                            if A_sim(a,i)>max(A)
                                A_sim(a,i) = max(A);
                            end
                            %W
                            if W_sim(a,i)>max(W)
                                W_sim(a,i) = max(W);
                            end	
                            %f                
                            if f_sim(a,i)<min(f) 
                                f_sim(a,i) = min(f);
                            elseif f_sim(a,i)>max(f)
                                f_sim(a,i) = max(f);
                            end        

                            %Enforce self-absorbing retirement (value functions impose it)
                            if  L_sim(a-1,i) == 1                 
                                L_sim(a,i) = round(lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Lfunc,A_sim(a,i),W_sim(a,i),f_sim(a,i)));                                     
                            elseif L_sim(a-1,i) == 0
                                L_sim(a,i) = 0;
                            end

                    end

                    LsimC{p,1} = L_sim;
                    AsimC{p,1} = A_sim;
                    WsimC{p,1} = W_sim;
                end

            %-------------------------------------------------------------------------%
            %Treated (i.e. DB frozen) group
            %-------------------------------------------------------------------------%

             for p = 1:9 
                L_sim = zeros(hor+6,obs(p));
                A_sim = zeros(hor+6,obs(p));
                W_sim = zeros(hor+6,obs(p));
                A_sim(1,:) = A0{p,1};
                W_sim(1,:) = W0{p,1};

                %Set disutility distribution
                f_sim = D{p,1};

                %Employment choice for t = -5 outside the loop (t = period
                %relative to freeze) a = 1 when t = -5
                for i = 1:obs(p)
                    a=1;    

                    %No-freeze policy functions
                    Afunc_w = optfunc{(aT-aS)-(p+a-1),2};
                    Wfunc_w = optfunc{(aT-aS)-(p+a-1),3};      
                    Afunc_r1 = optfunc{(aT-aS)-(p+a-1),5};
                    Afunc_r  = Afunc_r1{p+a};
                    Wfunc_r1 = optfunc{(aT-aS)-(p+a-1),6};  
                    Wfunc_r  = Wfunc_r1{p+a};                  
                    Lfunc = optfunc{(aT-aS)-(p+a-1),8};

                    %Put extreme draws of A,W,f onto the grid 
                    %A
                    if A_sim(a,i)>max(A)
                        A_sim(a,i) = max(A);
                    end
                    %W
                    if W_sim(a,i)>max(W)
                        W_sim(a,i) = max(W);
                    end	
                    %f                
                    if f_sim(a,i)<min(f) 
                        f_sim(a,i) = min(f);
                    elseif f_sim(a,i)>max(f)
                        f_sim(a,i) = max(f);
                    end      

                        %Simulated values 
                        L_sim(a,i) = round(lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Lfunc,A_sim(a,i),W_sim(a,i),f_sim(a,i)));
                        if L_sim(a,i) == 1
                            A_sim(a+1,i) = lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Afunc_w,A_sim(a,i),W_sim(a,i),f_sim(a,i));
                            W_sim(a+1,i) = lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Wfunc_w,A_sim(a,i),W_sim(a,i),f_sim(a,i));
                        elseif L_sim(a,i) == 0 
                            A_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Afunc_r,A_sim(a,i),W_sim(a,i));
                            W_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Wfunc_r,A_sim(a,i),W_sim(a,i));   
                        end  
                end

                for a = 2:5 % pre-freeze after initial period
                   for i = 1:obs(p)

                        %No-freeze policy functions
                        Afunc_w = optfunc{(aT-aS)-(p+a-1),2};
                        Wfunc_w = optfunc{(aT-aS)-(p+a-1),3};      
                        Afunc_r1 = optfunc{(aT-aS)-(p+a-1),5};
                        Afunc_r  = Afunc_r1{p+a};
                        Wfunc_r1 = optfunc{(aT-aS)-(p+a-1),6};  
                        Wfunc_r  = Wfunc_r1{p+a};                 
                        Lfunc = optfunc{(aT-aS)-(p+a-1),8};

                        %Put extreme draws of A,W,f onto the grid 
                        %A
                        if A_sim(a,i)>max(A)
                            A_sim(a,i) = max(A);
                        end
                        %W
                        if W_sim(a,i)>max(W)
                            W_sim(a,i) = max(W);
                        end	
                        %f                
                        if f_sim(a,i)<min(f) 
                            f_sim(a,i) = min(f);
                        elseif f_sim(a,i)>max(f)
                            f_sim(a,i) = max(f);
                        end        

                        %Enforce self-absorbing retirement (value functions impose it)
                        if  L_sim(a-1,i) == 1                 
                            L_sim(a,i) = round(lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Lfunc,A_sim(a,i),W_sim(a,i),f_sim(a,i)));                   
                                if L_sim(a,i) == 1
                                    A_sim(a+1,i) = lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Afunc_w,A_sim(a,i),W_sim(a,i),f_sim(a,i));
                                    W_sim(a+1,i) = lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Wfunc_w,A_sim(a,i),W_sim(a,i),f_sim(a,i));
                                elseif L_sim(a,i) == 0 
                                    A_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Afunc_r,A_sim(a,i),W_sim(a,i));
                                    W_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Wfunc_r,A_sim(a,i),W_sim(a,i));   
                                end                    
                        elseif L_sim(a-1,i) == 0
                            L_sim(a,i) = 0;
                            A_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Afunc_r,A_sim(a,i),W_sim(a,i));
                            W_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Wfunc_r,A_sim(a,i),W_sim(a,i));
                        end

                    end
                end


                for a = 6:hor+5 %after freeze
                   for i = 1:obs(p)

                        %Freeze policy functions
                        Afunc_wZ = optfunc_frz{(aT-aS)-(p+a-1),p,2};
                        Wfunc_wZ = optfunc_frz{(aT-aS)-(p+a-1),p,3};      
                        Afunc_rZ1 = optfunc_frz{(aT-aS)-(p+a-1),p,5};
                        Afunc_rZ  = Afunc_rZ1{p+a};
                        Wfunc_rZ1 = optfunc_frz{(aT-aS)-(p+a-1),p,6};  
                        Wfunc_rZ  = Wfunc_rZ1{p+a};
                        LfuncZ = optfunc_frz{(aT-aS)-(p+a-1),p,8};

                        %Put extreme draws of A,W,f onto the grid 
                        %A
                        if A_sim(a,i)>max(A)
                            A_sim(a,i) = max(A);
                        end
                        %W
                        if W_sim(a,i)>max(W)
                            W_sim(a,i) = max(W);
                        end	
                        %f                
                        if f_sim(a,i)<min(f) 
                            f_sim(a,i) = min(f);
                        elseif f_sim(a,i)>max(f)
                            f_sim(a,i) = max(f);
                        end        

                        %Enforce self-absorbing retirement (value functions impose it)
                        if  L_sim(a-1,i) == 1                 
                            L_sim(a,i) = round(lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,LfuncZ,A_sim(a,i),W_sim(a,i),f_sim(a,i)));                   
                                if L_sim(a,i) == 1
                                    A_sim(a+1,i) = lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Afunc_wZ,A_sim(a,i),W_sim(a,i),f_sim(a,i));
                                    W_sim(a+1,i) = lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,Wfunc_wZ,A_sim(a,i),W_sim(a,i),f_sim(a,i));
                                elseif L_sim(a,i) == 0 
                                    A_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Afunc_rZ,A_sim(a,i),W_sim(a,i));
                                    W_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Wfunc_rZ,A_sim(a,i),W_sim(a,i));   
                                end                    
                        elseif L_sim(a-1,i) == 0
                            L_sim(a,i) = 0;
                            A_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Afunc_rZ,A_sim(a,i),W_sim(a,i));
                            W_sim(a+1,i) = lininterp2(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),Wfunc_rZ,A_sim(a,i),W_sim(a,i));
                        end

                    end
                end

                %Final period employment choice (A,W already determined)
                a = hor+6;
                for i = 1:obs(p)

                    %Freeze policy function   
                    LfuncZ = optfunc_frz{(aT-aS)-(p+a-1),p,8};

                    %Put extreme draws of A,W,f onto the grid 
                    %A
                    if A_sim(a,i)>max(A)
                        A_sim(a,i) = max(A);
                    end
                    %W
                    if W_sim(a,i)>max(W)
                        W_sim(a,i) = max(W);
                    end	
                    %f                
                    if f_sim(a,i)<min(f) 
                        f_sim(a,i) = min(f);
                    elseif f_sim(a,i)>max(f)
                        f_sim(a,i) = max(f);
                    end        

                    %Enforce self-absorbing retirement (value functions impose it)
                    if  L_sim(a-1,i) == 1                 
                        L_sim(a,i) = round(lininterp3(A.*(1+r)^(p+a-1),W.*(1+r)^(p+a-1),f,LfuncZ,A_sim(a,i),W_sim(a,i),f_sim(a,i)));                                  
                    elseif L_sim(a-1,i) == 0
                        L_sim(a,i) = 0;
                    end

                end

                LsimT{p,1} = L_sim;
                AsimT{p,1} = A_sim;
                WsimT{p,1} = W_sim;

             end

            %-------------------------------------------------------------------------%
            % Compute moments by averaging over simulated individuals
            %-------------------------------------------------------------------------%

            %Concatenate control group       
            L_C2 = [LsimC{1,1} LsimC{2,1} LsimC{3,1} LsimC{4,1} LsimC{5,1} LsimC{6,1} ...
                    LsimC{7,1} LsimC{8,1} LsimC{9,1}];

            %Condition on sample of workers who are still in LF 5 years pre-freeze
            indx_1 = find(L_C2(1,:)==0);

            L_C2_r1 = L_C2;
            L_C2_r1(:,indx_1)=[];

            %Extract treated group data        
            L_T2 = [LsimT{1,1} LsimT{2,1} LsimT{3,1} LsimT{4,1} LsimT{5,1} LsimT{6,1} ...
                    LsimT{7,1} LsimT{8,1} LsimT{9,1}];

            indx_2 = find(L_T2(1,:)==0);

            L_T2_r1 = L_T2;
            L_T2_r1(:,indx_2)=[];

            %Make simulated data moments
            %ensure that the data matrix is not empty
            %1. Employment rates conditional on working 5 years pre-freeze
            if isempty(L_C2_r1)==0
                EL_C2r = mean(L_C2_r1(2:end,:),2); %mean employment rates
            else
                EL_C2r = ones(hor+5,1); 
            end    
            %2. employment rates and ATE of freeze conditional on working 5 yrs pre freeze
            %age
            if isempty(L_C2_r1)==0 && isempty(L_T2_r1)==0
                EL_C2 = mean(L_C2_r1(6:end,:),2); %mean ER control
                EL_T2 = mean(L_T2_r1(6:end,:),2); %mean ER affected
                TE_2= (EL_T2-EL_C2); %Teffect (pp  diff in ER's)
            else
                TE_2= zeros(hor,1);
            end

            %Extract and scale the simulated moments 
            E_Se = EL_C2r(:,1); 
            TE_2e = scaler.*TE_2;
            smom = [E_Se;TE_2e];

            %Write the data 
            Sm(row,:)=smom';
        end

    %Compute the VCV matrix: 
    %subtract mean of Sm from each row of Sm (residuals)
    %compute cross product of residuals and multiply by Nreps^-1

    W_matrix = (Nreps^-1).*(Sm - mean(Sm))'*(Sm-mean(Sm));

end

